from typing import Union, IO

import tomlkit
from tomlkit import TOMLDocument


class TomlParser(object):

    def __init__(self, doc: TOMLDocument):
        self.doc = doc

    def get(self, key_path: str, default=None):
        keys = key_path.split('.')
        tmp = self.doc
        length = len(keys)
        for index, key in enumerate(keys, 1):
            if isinstance(tmp, dict) and key in tmp:
                tmp = tmp[key]
                if index == length:
                    return tmp
            else:
                return default

    def set(self):
        raise NotImplemented()

    @classmethod
    def parse_from_string(cls, content: str) -> 'TomlParser':
        doc = tomlkit.parse(content)
        return cls(doc)

    @classmethod
    def parse_from_file(cls, file: Union[str, IO]) -> 'TomlParser':
        if hasattr(file, 'read'):
            content = file.read()
        else:
            with open(file) as f:
                content = f.read()
        return cls.parse_from_string(content)


if __name__ == '__main__':
    d = TomlParser.parse_from_file('./conf.toml')
    print(d.get('collect_gateways'))
    for x in d.get('collect_gateways'):
        print(x)

